# -*- coding: utf-8 -*-
# FileName:     zplot.py
# time:         23/5/4 004 下午 5:13
# Author:       Zhou Hang
# Description:  Hardware 使用PIL实现的绘图
from PIL import Image, ImageDraw
import numpy as np
import math

import Adafruit_SSD1306

from src.utils.common import RESOURCE_DIR, SCR_W, SCR_H, MIN16BIT, MAX16BIT


def data_remap(data: np.ndarray, a: int, b: int, min_: int, max_: int):
    """
    把一段范围在[min_, max_]之间的数字映射到[a, b]上
    这里a和b并没有谁大谁小的要求
    """
    delta = max_ - min_
    return np.round(a + (b - a) / delta * (data - min_))


class ZPlot:
    def __init__(self, w=SCR_W, h=SCR_H, mode="wave"):
        """
        :param mode: freq - 频谱模式 | wave - 波谱模式
        """
        self.w = w
        self.h = h
        self.mode = mode
        self.frame = Image.new("1", (w, h))
        self.draw = ImageDraw.Draw(self.frame)

        self.disp = Adafruit_SSD1306.SSD1306_128_64(rst=None)
        self.disp.begin()
        self.w = self.disp.width
        self.h = self.disp.height
        # self.disp.clear()
        # self.disp.display()

    def display(self, filepath: str):
        """
        将filepath指定的二值图放到OLED上
        """
        try:
            frame = Image.open(filepath)
        except FileNotFoundError:  # do nothing
            frame = Image.open(f"{RESOURCE_DIR}/error.jpg")
        frame = frame.convert("1")
        self.disp.image(frame)
        self.disp.display()
        frame.close()

    def set_mode(self, mode):
        self.mode = mode

    def plot(self, y):
        """
        传入的数组y长度不为0即可
        """
        if self.mode == "wave":
            # 将传入参数长度平均抽取为SCR_H
            y = y[::math.ceil(len(y) / self.w)]
            # 将数值映射到[0, h]上
            y = data_remap(y, 0, self.h - 1, MIN16BIT // 9, (MAX16BIT - 1) // 9)
            # y = data_remap(y, 0, self.h - 1, np.min(y), np.max(y))
            assert len(y) <= self.w
        elif self.mode == "freq":
            ffty = np.fft.fft(y)
            y = ffty / (len(y) / 2)  # 换算成实际的振幅
            y[0] = y[0] / 2
            y = abs(y[range(1, int(len(y) / 2))])
            y = y[::math.ceil(len(y) / self.w)]
            y = data_remap(y, self.h - 1, 0, 0, 50)

        for i in range(len(y)):
            if y[i] >= 64 or y[i] <= 0:  # 当数据超过设置的软上限时就会溢出，此时不做处理即可
                continue
            self.frame.putpixel((i, int(y[i])), 1)
            # elif VIS_MODE == "freq":

        self.disp.image(self.frame)
        self.disp.display()
        self.clear()

    def get_frame(self):
        return self.frame

    def show(self):
        self.frame.show()

    def clear(self):
        self.frame = Image.new("1", (self.w, self.h))
        # self.draw.rectangle((0, 0, self.w, self.h), outline=0, fill=0)


class ZPlotEmpty:
    """
    用于在i2c屏幕没有连接的时候do nothing
    """
    def __init__(self, w=SCR_W, h=SCR_H, mode="1") -> None:
        print("请检查i2c设备是否正确连接")
        pass

    def display(self, filepath: str):
        pass

    def plot(self, y):
        pass

    def get_frame(self):
        pass

    def show(self):
        pass

    def clear(self):
        pass


def main():
    plt = ZPlot(128, 64, '1')
    y = range(64)
    plt.plot(y)
    plt.show()


if __name__ == "__main__":
    main()
